import os
import cv2
import numpy as np
from shutil import copyfile

# 定义原始数据集路径和输出新数据集路径
data_dir = 'data/driving/train'
new_data_dir = 'data/driving/new_train'
# 统计原始的数据集中各类别图像数量
for i in range(10):
    subdir = os.path.join(data_dir, 'c{}'.format(i))
    num_files = len(os.listdir(subdir))
    print('old_Class {}: {} files'.format(i, num_files))
# 加载原始数据集中的所有图像
data_files = []
for i in range(10):
    subdir = os.path.join(data_dir, 'c{}'.format(i))
    for filename in os.listdir(subdir):
        filepath = os.path.join(subdir, filename)
        img = cv2.imread(filepath)
        data_files.append((img, i))
# 定义噪声数据集需要生成的数量和输出路径
num_noise_files = len(data_files) // 5
noise_dir = 'data/driving/noise'
# 创建输出文件夹
os.makedirs(new_data_dir, exist_ok=True)
os.makedirs(noise_dir, exist_ok=True)

# 定义噪声生成函数
def generate_noise_image(img):
    # 模拟高斯噪声
    noise = np.random.normal(size=img.shape, scale=50)
    noise_img = img + noise
    noise_img = np.clip(noise_img, 0, 255).astype(np.uint8)
    return noise_img


# 生成噪声数据并添加到新的数据集中
for i in range(num_noise_files):
    # 随机选择一个原始图像，并生成噪声图像
    img_idx = np.random.randint(len(data_files))
    img, label = data_files[img_idx]
    noise_img = generate_noise_image(img)
    # 保存噪声图像到文件夹中
    filename = 'noise_{:05d}.jpg'.format(i)
    filepath = os.path.join(noise_dir, filename)
    cv2.imwrite(filepath, noise_img)

    # 将噪声图像添加到新的数据集中
    new_subdir = os.path.join(new_data_dir, 'c{}'.format(label))
    os.makedirs(new_subdir, exist_ok=True)
    new_filename = 'noise_{:05d}.jpg'.format(i)
    new_filepath = os.path.join(new_subdir, new_filename)
    copyfile(filepath, new_filepath)
# 将原始数据集中的图像复制到新的数据集中
for i in range(10):
    subdir = os.path.join(data_dir, 'c{}'.format(i))
    new_subdir = os.path.join(new_data_dir, 'c{}'.format(i))
    os.makedirs(new_subdir, exist_ok=True)
    for filename in os.listdir(subdir):
        filepath = os.path.join(subdir, filename)
        new_filepath = os.path.join(new_subdir, filename)
        copyfile(filepath, new_filepath)
# 统计新的数据集中各类别图像数量
for i in range(10):
    subdir = os.path.join(new_data_dir, 'c{}'.format(i))
    num_files = len(os.listdir(subdir))
    print('Class {}: {} files'.format(i, num_files))
